import io
import os
import tempfile
from urllib.parse import quote

import uvicorn
from fastapi import FastAPI, UploadFile
from starlette.responses import RedirectResponse, FileResponse, StreamingResponse

from marker.convert import convert_single_pdf
from marker.models import load_all_models


async def document():
    return RedirectResponse(url="/docs")


def mount_app_routes(app):
    app.get("/",
            summary="swagger文档")(document)


app = FastAPI(
    title="marker API Server",
    version="v1.0"
)
mount_app_routes(app)
model_lst = []


@app.post("/convert_single", summary="pdf文件转换", )
async def convert_single(file: UploadFile, max_pages: int = None, ):
    # 读取上传文件的内容
    file_content = await file.read()
    # 创建临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        temp_file.write(file_content)
        temp_file_path = temp_file.name  # 获取临时文件路径

    full_text, _, out_meta = convert_single_pdf(temp_file_path, model_lst, max_pages=max_pages,
                                                batch_multiplier=2)
    file_data = io.BytesIO(full_text.encode('utf-8'))
    os.remove(temp_file_path)
    filename = file.filename.rsplit('.', 1)[0]
    # 使用 quote 对文件名进行 URL 编码
    quoted_filename = quote(f"{filename}.md")
    # 返回文件内容
    return StreamingResponse(file_data, media_type="application/octet-stream", headers={
        "Content-Disposition": f"attachment; filename={quoted_filename}"
    })


if __name__ == '__main__':
    model_lst = load_all_models()
    uvicorn.run(app, host="0.0.0.0", port=8862)
